import oneflow as flow
import logging
import oneflow.nn as nn
from oasr.model.base import BaseModel
from oasr.encoder import Encoder
from oasr.decoder import Decoder
from oasr.frontend import FrontEnd
from oasr.module.loss import LabelSmoothingLoss
from oasr.model.ctc import CTCAssistor


class SpeechToText(BaseModel):
    def __init__(self, params):
        super(SpeechToText, self).__init__()

        if 'frontend' in params:
            self.include_frontend = True 
            self.frontend = FrontEnd[params['frontend_type']](**params['frontend'])
        else:
            self.include_frontend = False

        self.encoder = Encoder[params['encoder_type']](**params['encoder'])
        logging.info('Build a %s encoder!' % params['encoder_type'])
        self.decoder = Decoder[params['decoder_type']](**params['decoder'])
        logging.info('Build a %s decoder!' % params['decoder_type'])

        self.crit = LabelSmoothingLoss(
            size=params['decoder']['vocab_size'],
            smoothing=params['smoothing']
        )

        self.ctc_weight = params['ctc_weight']
        if self.ctc_weight > 0.0:
            self.assistor = CTCAssistor(
                hidden_size=params['encoder_output_size'],
                vocab_size=params['decoder']['vocab_size'],
                lookahead_steps=params['lookahead_steps'] if 'lookahead_steps' in params else 0)
            logging.info('Build a CTC Assistor with weight %.2f' % self.ctc_weight)
        
    def forward(self, inputs, targets):

        enc_inputs = inputs['inputs']
        enc_mask = inputs['mask']

        truth = targets['targets']
        truth_length = targets['targets_length']

        if self.include_frontend:
            enc_inputs, enc_mask = self.frontend(enc_inputs, enc_mask)

        # 1. forward encoder
        memory, memory_mask, _ = self.encoder(enc_inputs, enc_mask)

        # 2. forward decoder
        target_in = truth[:, :-1].clone()
        logits, _ = self.decoder(target_in, memory, memory_mask)

        # 3. compute attention loss
        target_out = truth[:, 1:].clone()
        loss = self.crit(logits, target_out)

        if self.ctc_weight > 0:
            loss_ctc = self.compute_ctc_loss(memory, memory_mask, target_out, truth_length)
            return (1 - self.ctc_weight) * loss + self.ctc_weight * loss_ctc, {'CTCLoss': loss_ctc.item()}
        else:
            return loss, None

    def compute_ctc_loss(self, memory, memory_mask, targets_out, targets_length):
        memory_length = flow.sum(memory_mask.squeeze(1), dim=-1)
        loss_ctc = self.assistor(memory, memory_length, targets_out, targets_length)
        return loss_ctc

    def save_checkpoint(self, params, name):
        checkpoint = {
            'params': params,
            'encoder': self.encoder.state_dict(),
            'decoder': self.decoder.state_dict()
            }

        if self.ctc_weight > 0.0:
            checkpoint['ctc'] = self.assistor.state_dict()

        if self.include_frontend:
            checkpoint['frontend'] = self.frontend.state_dict()

        flow.save(checkpoint, name)

    def set_epoch(self, epoch):
        pass
